# provide translation
# author-by: xjtu-blacksmith
# create-on: 2020.2.22

import os
from os import path
from googletrans import Translator
from progressbar import ProgressBar

CATEGORIES = path.join('data', 'categories.txt')  # pre-produced category list
DATASET_DIR = path.join('data', 'caltech-101')
REFERENCE_SET = path.join(DATASET_DIR, 'train')
TRANS_DICT = path.join(DATASET_DIR, 'translation.txt')

def build_translation(input_dir=REFERENCE_SET, output=TRANS_DICT):
    ''' build translation dict by googletrans
    '''
    with open(CATEGORIES, 'r') as f:
        categories = f.readlines()
        if categories[-1] == '':  # remove empty line in the end
            categories.pop()
    lines = []  # lines to be written in dict file
    translator = Translator()  # google translator
    bar = ProgressBar()  # default progressbar
    for i in bar(range(len(categories))):
        category = categories[i].replace('_', ' ').rstrip('\n')  # replace underlines
        line = translator.translate(category, dest='zh-CN').text
        lines.append(line + '\n')  # newline at end
    with open(output, 'w', encoding='utf-8') as f:
        f.writelines(lines)
        print('Translation done, %d names translated.' % len(lines))
        print('翻译完成，共翻译了 %d 个名称。' % len(lines))

def get_class(nid):
    with open(CATEGORIES, 'r') as f:
        categories = f.readlines()
        return categories[nid].rstrip('\n')

def translate_id(nid, dict_dir=None, encode='utf-8'):
    ''' translate the order number to Chinese meaning
    '''
    if dict_dir is None:
        dict_dir = TRANS_DICT
    with open(dict_dir, 'r', encoding='utf-8') as f:
        lines = f.readlines() # read all lines
        meaning = lines[nid].rstrip('\n')  # get the meaning
    if encode != 'utf-8':
        meaning = meaning.encode(encoding=encode)
    return meaning

if __name__ == "__main__":
    if path.exists(TRANS_DICT) == False:
        build_translation()
    print(translate_id(0))